﻿#ifndef __BST__
#define __BST__
#include <iostream>
#include <stack>

#define RECURSION 1
#define ROOT_INSERT 1

template <class item>
class BST{
public:
	//外部接口
	BST();
	~BST();
	bool empty() const;
	bool search(item);
	void insert(item);
	item select(int);//选择第N小的项目
	void traverse();
	void remove(item);
	void part(int n) {part_in(n, root);}
private:
	//内部数据结构
	struct node{
		item data;
		node *left, *right;
		int child_count;// 子节点的个数
		node(item& data):data(data),left(NULL),right(NULL),child_count(0){}
	};
	typedef node* link;
	link root;
	//内部实现的接口
	bool search_in(item, link);
	void insert_in(item, link&);
	void remove_in(item, link&);
	void part_in(int, link&);//将第N小的项目作为根，达到分割左右两树的目的
	item select_in(int, link);
	void left_rotate(link&);
	void left_rotate(link&,link&);
	void right_rotate(link&);
	void right_rotate(link&,link&);
	link join_left_right(link left, link right){
		if(right == NULL) return left;
		part_in(1, right);
		right->left = left;
		right->child_count += left ? left->child_count + 1 : 0
		return right;
	}
	void traverse_in(std::ostream&, link);
	void destroy(link&);
};


#if RECURSION
template <class item>
void BST<item>::destroy(link& root){
	if(root == NULL) return;
	if(root->left != NULL) destroy(root->left);
	if(root->right != NULL) destroy(root->right);
	delete root; root = NULL;
}

template <class item>
bool BST<item>::search_in(item key, link root){
	if(root == NULL) return false;
	if(key == root->data) return true;
	else if(key > root->data) return search_in(key, root->right);
	else return search_in(key, root->left);
}

template <class item>
item BST<item>::select_in(int n, link root){
	assert(root != NULL && n > 0);
	int left_child_cnt = (root->left == NULL)? 0 : root->left->child_count+1;
	if(n == left_child_cnt+1) return root->data;
	if(n < left_child_cnt+1) return select_in(n, root->left);
	else return select_in(n-left_child_cnt-1, root->right);
}

template <class item>
void BST<item>::part_in(int n, link& root){
	if(root == NULL || n < 1) return;
	int lcc = root->left ? root->left->child_count + 1 : 0;
	if(lcc+1 == n) return;
	if( n < lcc+1) { part_in(n, root->left); right_rotate(root); }
	else { 	part_in(n-lcc-1, root->right); left_rotate(root);	}
}

//为了维护child_count字段，我们在删除前判断是否存在该项
//否则应该在删除失败的情况下，进行child_count 字段的回滚操作
template <class item>
void BST<item>::remove_in(item key, link& root){
	if(root == NULL || !search_in(key, root)) return;
	if(key == root->data){
		link del = root;
		root = join_left_right(root->left, root->right);
		delete del;
	}else if(key > root->data) { remove_in(key,  root->right); --root->child_count; }
	else { remove_in(key, root->left); --root->child_count;}
}

template <class item>
void BST<item>::right_rotate(link& node){
	//旋转后child_count 字段的变化
	int add = node->right ? node->right->child_count + 1 + 1: 1 ;
	int del = node->left ? (node->left->left ? node->left->left->child_count + 1 + 1 : 1) : 1;
	node->child_count -= del; 	node->left->child_count += add;
	//旋转
	link x = node->left; node->left = x->right; x->right = node; node = x;
}

template <class item>
void BST<item>::left_rotate(link& node){
	//旋转后child_count 字段的变化
	int add = node->left ? node->left->child_count + 1 + 1 : 1 ;
	int del = node->right ? (node->right->right ? node->right->right->child_count + 1 + 1 : 1) : 1;
	node->child_count -= del; 	node->right->child_count += add;
	//旋转
	link x = node->right; node->right = x->left; x->left = node; node = x;
}

#if ROOT_INSERT//将插入项，旋转为根
template <class item>
void BST<item>::insert_in(item key, link& root){
	if(search_in(key, root)) return;
	if(root == NULL) { root = new node(key); return; }
	if(key < root->data) {	insert_in(key, root->left); ++root->child_count; right_rotate(root);	}
	else if(key > root->data){ 	insert_in(key, root->right); ++root->child_count; left_rotate(root);	}
	else return;
}
#else
template <class item>
void BST<item>::insert_in(item key, link& root){
	if(search_in(key, root)) return;
	if(root == NULL) { root = new node(key); return; }
	if(key < root->data) { insert_in(key, root->left); ++root->child_count;}
	else { insert_in(key, root->right); ++root->child_count; }
}

#endif//根插入 选择

template <class item>//前序遍历
void BST<item>::traverse_in(std::ostream& out, link root){
	if(root == NULL) return;
	out << root->data << " ";
	traverse_in(out, root->left);
	traverse_in(out, root->right);
}


#else
template <class item>
void BST<item>::destroy(link& root){
	if(root == NULL) return;
	std::stack<link> link_stack;
	link_stack.push(root);
	while(!link_stack.empty()){
		link top = link_stack.top(); link_stack.pop();
		if(top->left != NULL) link_stack.push(top->left);
		if(top->right != NULL) link_stack.push(top->right);
		delete top;
	}
}

template <class item>
bool BST<item>::search_in(item key, link root){
	if(root == NULL) return false;
	while(root != NULL){
		if(key == root->data) return true;
		else if(key < root->data) root = root->left;
		else root = root->right;
	}
	return false;
} 

#if ROOT_INSERT
template <class item>
void BST<item>::insert_in(item key, link& root){
	link father = NULL, cursor = root;
	std::stack<link> father_stack;
	std::stack<link> grandpa_stack;
	std::stack<int> dir_stack;
	while(cursor != NULL){
		if(key == cursor->data) return;
		else{
			if(!father_stack.empty())
				grandpa_stack.push(father_stack.top());
			father = cursor;
			father_stack.push(father);
			if(key < cursor->data)  { cursor = cursor->left; dir_stack.push(0); } 
			else { cursor = cursor->right; dir_stack.push(1); }
		}
	}
	if(father == NULL) { root = new node(key); return; }
	if(key < father->data) father->left = new node(key);
	else father->right = new node(key);
	while(!father_stack.empty() && !dir_stack.empty()){
		link top = father_stack.top();
		++top->child_count;
		link parent = grandpa_stack.empty()? NULL: grandpa_stack.top();
		int dir = dir_stack.top();
		if(dir == 0) right_rotate(top, parent);
		else left_rotate(top, parent);
		if(father_stack.size() == 1)
			this->root = top;
		father_stack.pop(); dir_stack.pop();
		if(!grandpa_stack.empty()) grandpa_stack.pop();
	}
}
#else
template <class item>
void BST<item>::insert_in(item key, link& root){
	link father = NULL, cursor = root;
	std::stack<link> father_stack;
	while(cursor != NULL){
		if(key == cursor->data) return;
		else{
			father = cursor;
			father_stack.push(father);
			if(key < cursor->data)  cursor = cursor->left; 
			else cursor = cursor->right;
		}
	}
	if(father == NULL) { root = new node(key); return; }
	if(key < father->data) father->left = new node(key);
	else father->right = new node(key);

	while(!father_stack.empty()){
		link top = father_stack.top();
		++top->child_count;
		father_stack.pop();
	}
}
#endif

template <class item>
item BST<item>::select_in(int n, link root){
	assert(root != NULL && n > 0);
	link father = NULL, cursor = root;
	while(cursor != NULL){
		int left_child_count = cursor->left? cursor->left->child_count + 1 : 0;
		if(n == left_child_count+1) return cursor->data;
		else{
			father = cursor;
			if(n < left_child_count+1) cursor = cursor->left;
			else { cursor = cursor->right; n = n-left_child_count-1; }
		}
	}
	assert(cursor != NULL);//循环后，cursor 为NULL则说明 n 越界了
}

template <class item>
void BST<item>::traverse_in(std::ostream& out, link root){
	if(root == NULL) return;
	std::stack<link> link_stack;
	link_stack.push(root);
	while(!link_stack.empty()){
		link top = link_stack.top(); link_stack.pop();
		if(top->left != NULL) link_stack.push(top->left);
		if(top->right != NULL) link_stack.push(top->right);
		out << top->data <<" ";
	}
}

template <class item>
void BST<item>::remove_in(item key, link& root){
	if(root == NULL ) return;
	link cursor = root , father = NULL;
	std::stack<link> link_stack;
	int dir = 0;
	while(cursor != NULL){
		if(key == cursor->data) break;
		else {
			father = cursor;
			link_stack.push(father);
			if(key < cursor->data) { cursor = cursor->left; dir = 0; }
			else { cursor = cursor->right; dir = 1; }
		}
	}
	if(cursor == NULL) return;
	link del = cursor;
	if(father == NULL) { this->root = join_left_right(del->left, del->right); }
	else {
		if(dir == 0)
			father->left = join_left_right(del->left, del->right);
		else father->right = join_left_right(del->left, del->right);
	}
	while(!link_stack.empty()){
		link top = link_stack.top(); link_stack.pop();
		--top->child_count;
	}
	delete del;
}

template <class item>
void BST<item>::part_in(int n, link& root){
	if(root == NULL || n < 1) return;
	link cursor = root;
	std::stack<link> father_stack;
	std::stack<link> grandpa_stack;
	std::stack<int> dir_stack;
	while(cursor != NULL){
		int llc = cursor->left? cursor->left->child_count + 1 : 0;
		if(n == llc+1) break ;
		else{
			if(!father_stack.empty())
				grandpa_stack.push(father_stack.top());
			father_stack.push(cursor);
			
			if(n < llc+1) { 
				dir_stack.push(0);
				cursor = cursor->left;				
			} else {
				dir_stack.push(1);
				cursor = cursor->right;
				n = n - llc - 1;
			}
		}
	}

	while(!father_stack.empty() && !dir_stack.empty()){
		link top = father_stack.top();
		link parent = grandpa_stack.empty()? NULL: grandpa_stack.top();
		int dir = dir_stack.top();
		if(dir == 0) right_rotate(top, parent);	
		else left_rotate(top, parent);
		if(father_stack.size() == 1)
			this->root = top;
		father_stack.pop(); dir_stack.pop();
		if(!grandpa_stack.empty()) grandpa_stack.pop();
	}
}

template <class item>
void BST<item>::right_rotate(link& node, link& parent){
	link* ptr = &node;
	int root_add = node->right ? node->right->child_count + 1 + 1 : 1;
	link x = node->left; node->left = x->right; x->right = node; 
	int father_del = x->left ? x->left->child_count + 1 + 1 : 1;
	x->child_count += root_add;
	node->child_count -= father_del;
	if(parent == NULL) { node = x; return;} 
	if(parent->left == *ptr) parent->left = x;
	else parent->right = x;
}

template <class item>
void BST<item>::left_rotate(link& node, link& parent){
	link* ptr = &node;
	int root_add = node->left ? node->left->child_count + 1 + 1 : 1;
	link x = node->right; node->right = x->left; x->left = node; 
	int father_del = x->right ? x->right->child_count + 1 + 1 : 1;
	x->child_count += root_add;
	node->child_count -= father_del;
	if(parent == NULL)  { node = x; return; } 
	if(parent->left == *ptr) parent->left = x;
	else parent->right = x;
}
#endif//RECURSION 版本选择


/*********************************************
               外部接口
********************************************/
template <class item>
BST<item>::BST(){ root = NULL; }

template <class item>
BST<item>::~BST(){ 	destroy(root);	root = NULL; }

template <class item>
bool BST<item>::empty() const { return root == NULL; }

template <class item>
void BST<item>::insert(item key){ insert_in(key, root); }

template <class item>
bool BST<item>::search(item key){ return search_in(key, root); }

template <class item>
void BST<item>::traverse(){ traverse_in(std::cout, root); }

template <class item>
void BST<item>::remove(item key){ remove_in(key, root); }

template <class item>
item BST<item>::select(int n){ return select_in(n, root); }


#endif